package version4.Server.rateLimit.impl;

import version4.Server.rateLimit.RateLimit;

public class TokenBucketRateLimit implements RateLimit {
    // RATE 和 CAPACITY 是常量，但是不在此处定义死
    // 我们希望这个速率和容量可以根据不同的服务进行调整
    private static int RATE; // 单位：ms，每多少ms生成一个令牌
    private static int CAPACITY;
    private volatile int curCapacity;
    private volatile long timeStamp = System.currentTimeMillis();

    public TokenBucketRateLimit(int rate, int capacity) {
        RATE = rate;
        CAPACITY = capacity;
        // 初始化容量
        curCapacity = CAPACITY;
    }
    @Override
    public synchronized boolean getToken() {
        if(curCapacity > 0) {
            // 如果容量大于0，就可以获取令牌
            curCapacity--;
            return true;
        } else {
            // 计算当前时间和上次时间的时间差
            long now = System.currentTimeMillis();
            long gap = now - timeStamp;
            // 如果时间差大于RATE，就可以生成令牌
            // 否则进行限流
            if(gap >= RATE) {
                curCapacity += (int)gap / RATE -1;
                curCapacity = Math.min(curCapacity, CAPACITY);
                timeStamp = now;
                return true;
            }
        }
        return false;
    }
}
